"""
    参考文档：
    https://milvus.io/docs/install_standalone-docker.md
    https://github.com/shibing624/text2vec

"""

import os
from typing import List

import numpy as np
import pandas as pd
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
from text2vec import SentenceModel

try:
    from conf.config import config, BASE_DIR
except ModuleNotFoundError:
    from conf.config import config, BASE_DIR


embedding_model = SentenceModel(model_name_or_path=os.path.join(BASE_DIR, r"embedding_models\text2vec_base_chinese"),
                                max_seq_length=128, device="cpu")


def get_embeddings(texts: List[str]) -> List:
    """
    获取文本的向量值(归一化)
    :param texts: 要向量化的文本列表
    :return: 向量化后的列表
    """
    global embedding_model
    embeddings = embedding_model.encode(sentences=texts, batch_size=64, show_progress_bar=True).tolist()

    # 归一化处理
    embeddings_normalized = list()
    for embedding in embeddings:
        # 计算向量的L2范数
        norm = np.linalg.norm(embedding)
        # 将向量除以它的L2范数，得到一个范数为1的向量
        embeddings_normalized.append((embedding / norm).tolist())

    return embeddings_normalized


def data_to_embeddings():
    """
    数据向量化
    :return:
    """
    # 读取区域id，城市id，区域名
    csv_path = os.path.join(BASE_DIR, 'docs/district_data.csv')
    csv_output_path = os.path.join(BASE_DIR, 'docs/district_data_within_vector.csv')
    df = pd.read_csv(filepath_or_buffer=csv_path, usecols=[0, 1, 2],
                     converters={'district_id': lambda x: int(x), 'city_id': lambda x: int(x)})

    # 将区域名向量化
    district_name_vectors = get_embeddings(df['district_name'].tolist())

    # 新增一列
    df = df.assign(district_name_vector=[str(vector) for vector in district_name_vectors])
    # 保存
    df.to_csv(path_or_buf=csv_output_path, index=False)


def create_collection(collection_name: str, dim: int):
    """
    创建集合
    :param collection_name: 集合名
    :param dim: 矢量维度
    :return:
    """
    try:
        connections.connect(
            alias=config.get("milvus").get("debug").get("alias"),
            host=config.get("milvus").get("debug").get("host"), port=config.get("milvus").get("debug").get("port"),
            user=config.get("milvus").get("debug").get("user"), password=config.get("milvus").get("debug").get("passwd")
        )

        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)

        fields = [
            FieldSchema(name="district_id", dtype=DataType.INT64, is_primary=True, auto_id=False, description="区域id"),
            FieldSchema(name="city_id", dtype=DataType.INT64, description="城市id"),
            FieldSchema(name="district_name_vector", dtype=DataType.FLOAT_VECTOR, dim=dim, description="区域名对应的向量"),
        ]
        schema = CollectionSchema(fields=fields, description='区域表')
        collection = Collection(name=collection_name, schema=schema, using='default')

        index_params = {
            'metric_type': "IP",
            'index_type': "FLAT",
            'params': {}
        }
        collection.create_index(field_name='district_name_vector', index_params=index_params)
    finally:
        if "collection" in dir():
            collection.release()
        if "connections" in dir():
            connections.disconnect("default")


def insert_data_to_collection():
    # 读取区域id，城市id，区域名向量
    csv_path = os.path.join(BASE_DIR, 'docs/district_data_within_vector.csv')
    df = pd.read_csv(filepath_or_buffer=csv_path, usecols=[0, 1, 3],
                     converters={'district_id': lambda x: int(x), 'city_id': lambda x: int(x),
                                 'district_name_vector': lambda x: eval(x)})

    # 写入向量数据库
    data = df.transpose().values.tolist()
    try:
        connections.connect(
            alias=config.get("milvus").get("debug").get("alias"),
            host=config.get("milvus").get("debug").get("host"), port=config.get("milvus").get("debug").get("port"),
            user=config.get("milvus").get("debug").get("user"), password=config.get("milvus").get("debug").get("passwd")
        )

        collection = Collection("district")
        collection.insert(data)
        collection.flush()
    finally:
        if "collection" in dir():
            collection.release()
        if "connections" in dir():
            connections.disconnect("default")


def delete_data(collection_name: str, expr: str):
    """
    根据布尔表达式删除数据
    :param collection_name: 集合名
    :param expr: 布尔表达式，比如 "id == 1"、 "id in [2, 4]"等
    :return:
    """
    try:
        connections.connect(
            alias=config.get("milvus").get("debug").get("alias"),
            host=config.get("milvus").get("debug").get("host"), port=config.get("milvus").get("debug").get("port"),
            user=config.get("milvus").get("debug").get("user"), password=config.get("milvus").get("debug").get("passwd")
        )

        collection = Collection(collection_name)
        expr = expr
        collection.delete(expr)
        collection.flush()
    finally:
        if "collection" in dir():
            collection.release()
        if "connections" in dir():
            connections.disconnect("default")


def get_ids(text: str, collection_name: str, anns_field: str, top_k: int = 3, expr=None) -> List[int]:
    """
    获取数据项的id列表
    :param text: 要查找的文本
    :param collection_name: 集合名
    :param anns_field: 用于搜索集合的矢量字段的名称
    :param top_k: 返回前几个
    :param expr: 标量查询的布尔表达式
    :return:
    """
    try:
        connections.connect(
            alias=config.get("milvus").get("debug").get("alias"),
            host=config.get("milvus").get("debug").get("host"), port=config.get("milvus").get("debug").get("port"),
            user=config.get("milvus").get("debug").get("user"), password=config.get("milvus").get("debug").get("passwd")
        )

        collection = Collection(collection_name)
        collection.load()

        vector = get_embeddings([text])[0]
        results = collection.search(
            data=[vector],
            anns_field=anns_field,
            param={"metric_type": "IP"},
            limit=top_k,
            expr=expr
        )

        ids = results[0].ids
        return ids
    finally:
        if "collection" in dir():
            collection.release()
        if "connections" in dir():
            connections.disconnect("default")


def main():
    # # 步骤1 创建数据集合和索引
    # create_collection(collection_name="district", dim=768)

    # # 步骤2 数据向量化
    # data_to_embeddings()

    # # 步骤3 插入数据库
    # insert_data_to_collection()

    # 步骤4 查找向量数据
    # 测试一
    # 1.不限定城市
    ls = get_ids(text="南山", collection_name="district",
                 anns_field="district_name_vector", top_k=3,
                 expr=None)
    print(ls)  # [8400, 26208, 30473] → 深圳市南山, 玉山市山南镇, 鹤岗南山区
    # 2.限定城市
    city_id = 544  # 鹤岗市
    ls = get_ids(text="南山", collection_name="district", anns_field="district_name_vector",
                 top_k=3, expr=f"city_id == {city_id}")
    print(ls)  # [30473, 30479, 30477] → 鹤岗南山区, 鹤岗兴山区, 鹤岗东山区

    # 测试二
    # 1.不限定城市
    ls = get_ids(text="潮阳区", collection_name="district",
                 anns_field="district_name_vector", top_k=3,
                 expr=None)
    print(ls)  # [24032, 14807, 15239] → 潮州市潮安区, 汕头市潮阳, 汕头潮南
    # 2.限定城市
    city_id = 71  # 汕头市
    ls = get_ids(text="潮阳区", collection_name="district", anns_field="district_name_vector",
                 top_k=3, expr=f"city_id == {city_id}")
    print(ls)  # [14807, 15239, 14808] → 汕头市潮阳, 汕头潮南, 汕头南澳


if __name__ == '__main__':
    main()
